######################################################################################################################
# Copyright (c) 2006-2009 Gabor Dorka (gabor.dorka@gmail.com)
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation
# files (the "Software"), to deal in the Software without
# restriction, including without limitation the rights to use,
# copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following
# conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.


######################################################################################################################
import sys
import os
import stat
import time
import math
import json
import re
import subprocess


######################################################################################################################
class misp_phase_t:


	######################################################################################################
	def __init__(self, id='', configs=[]):
		self.id = id
		self.newly_generated = False
		self.configs = configs
		self.formulas = {}
		self.id_printed = False


	######################################################################################################
	def to_json_data(self):
		data = {}
		data['id'] = self.id
		data['formulas'] = formulas_data = {}
		for id, formula in self.formulas.items():
			formulas_data[id] = formula.to_json_data()
		return data


######################################################################################################################
class misp_formula_t:


	######################################################################################################
	def __init__(self, id, sources=[], targets=[], actions=[], extra_sources_name='', extra_targets_name=''):
		self.id = id
		self.newly_generated = True
		self.sources = sources
		self.targets = targets
		self.actions = actions
		self.extra_sources_name = extra_sources_name
		self.extra_targets_name = extra_targets_name
		self.sources_fileinfo = []
		self.extra_sources = []
		self.extra_sources_fileinfo = []
		self.targets_fileinfo = []
		self.extra_targets = []
		self.extra_targets_fileinfo = []
		self.prev_actions = []
		self.state_valid = True
		self.state_present = False
		self.phase = None


	######################################################################################################
	def to_json_data(self):
		data = {}
		data['newly_generated'] = self.newly_generated
		attr_names = ['sources', 'targets', 'actions', 'extra_sources_name', 'extra_targets_name']
		if self.state_valid:
			attr_names += ['extra_sources', 'extra_targets']
		for attr_name in attr_names:
			if len(self.__dict__[attr_name]) > 0:
				data[attr_name] = self.__dict__[attr_name]
		if self.state_valid:
			for attr_name in ['sources_fileinfo', 'extra_sources_fileinfo', 'targets_fileinfo', 'extra_targets_fileinfo']:
				if len(self.__dict__[attr_name]) > 0:
					data[attr_name] = attr_data = []
					for attr in self.__dict__[attr_name]:
						attr_data.append(attr.to_json_data())
		return data


	######################################################################################################
	def from_json_data(self, data):
		attr_names = ['sources', 'targets', 'actions', 'extra_sources_name', 'extra_sources', 'extra_targets_name', 'extra_targets']
		for attr_name in attr_names:
			attr = data.get(attr_name)
			if attr is not None:
				self.__dict__[attr_name] = attr
		for attr_name in ['sources_fileinfo', 'extra_sources_fileinfo', 'targets_fileinfo', 'extra_targets_fileinfo']:
			attr = data.get(attr_name)
			if attr is not None:
				for name, mtime in attr:
					self.__dict__[attr_name].append(misp_fileinfo_t(name, mtime))


	######################################################################################################
	def copy_state(self, prev):
		self.state_present = True
		for attr_name in ['extra_sources', 'sources_fileinfo', 'extra_sources_fileinfo', 'targets_fileinfo', \
							'extra_targets', 'extra_targets_fileinfo']:
			self.__dict__[attr_name] = prev.__dict__[attr_name]
		self.prev_actions = prev.actions


######################################################################################################################
class misp_fileinfo_t:


	######################################################################################################
	def __init__(self, name, mtime, exists=False, is_dir=False):
		self.name = name
		self.mtime = mtime
		self.exists = exists
		self.is_dir = is_dir


	######################################################################################################
	def to_json_data(self):
		return [self.name, self.mtime]


######################################################################################################################
class misp_t:


	######################################################################################################
	VERSION = '1.08'
	JSON_INDENT = 4
	MUST_EXIST = True
	DONT_HAVE_TO_EXIST = False
	USE_CACHE = True
	DONT_USE_CACHE = False
	CHECK_EQUAL = True
	CHECK_OLDER = False
	DISPLAY_HEADER_STRS = set(['action', 'always', 'never'])


	######################################################################################################
	def __init__(self):
		self.param_str = ''
		self.param_data = None
		self.mispgen_action = None
		self.misptemp_name = u'misptemp'
		self.mispfile_name = u'mispfile'
		self.exec_config = u''
		self.exec_phase_ids = []
		self.exec_formula = u''
		self.options = {}
		self.multi_pass_gen = False
		self.configs = {}
		self.prev_phases = []
		self.phases = []
		self.current_phase = None
		self.formulas = {}
		self.fileinfo_map = {}
		self.display_mispgen = True
		self.display_phase = u'action'
		self.display_formula = u'action'
		self.phase_format = u'=== %s ==='
		self.formula_format = u'--- %s ---'


	######################################################################################################
	def main(self):
		start_time = time.clock()
		sys.stdout.write('misp ' + self.VERSION + '\n')
		self.read_cmd_line()
		self.read_mispfile()
		self.generate_misptemp('')
		try:
			self.execute_phase_list()
		except KeyboardInterrupt:
			pass
		self.write_mispfile()
		end_time = time.clock()
		exec_secs = end_time - start_time
		exec_minutes = math.floor(exec_secs / 60.0)
		exec_secs -= exec_minutes * 60.0
		sys.stdout.write('\nexec time: %02.0f:%05.2f secs\n' % (exec_minutes, exec_secs))


	######################################################################################################
	def read_cmd_line(self):
		if len(sys.argv) != 2:
			self.error_exit('there must be exactly 1 param in JSON format')
		self.param_str = sys.argv[1]
		try:
			self.param_data = json.loads(self.param_str)
		except Exception as exc:
			self.error_exit('JSON load error: '+str(exc))
		self.mispgen_action = self.param_data.get('mispgen_action')
		if self.mispgen_action is None:
			self.error_exit("mispgen_action mustn't be empty")
		self.exec_config = self.param_data.get('exec_config')
		if self.exec_config is None or type(self.exec_config)!=type(u'') or len(self.exec_config)==0:
			self.error_exit("exec_config must be a non-empty string")
		self.exec_formula = self.param_data.get('exec_formula', self.exec_formula)
		if type(self.exec_formula) != type(u''):
			self.error_exit("exec_formula must be a string")
		self.misptemp_name = self.param_data.get('misptemp_name', self.misptemp_name)
		if type(self.misptemp_name) != type(u'') or len(self.misptemp_name) == 0:
			self.error_exit("misptemp_name must be a non-empty string")
		mispfile_name = self.param_data.get('mispfile_name', self.mispfile_name)
		if type(mispfile_name) != type(u'') or len(mispfile_name) == 0:
			self.error_exit("mispfile_name must be a non-empty string")
		self.mispfile_name = mispfile_name
		self.display_mispgen = self.param_data.get('display_mispgen', self.display_mispgen)
		if type(self.display_mispgen) != type(True):
			self.error_exit("display_mispgen must be a boolean")
		self.display_phase = self.param_data.get('display_phase', self.display_phase)
		if self.display_phase not in self.DISPLAY_HEADER_STRS:
			self.error_exit("display_phase value must be in %s" % self.DISPLAY_HEADER_STRS)
		self.display_formula = self.param_data.get('display_formula', self.display_formula)
		if self.display_formula not in self.DISPLAY_HEADER_STRS:
			self.error_exit("display_formula value must be in %s" % self.DISPLAY_HEADER_STRS)
		self.phase_format = self.param_data.get('phase_format', self.phase_format)
		try:
			s = self.phase_format % u'str'
		except TypeError:
			self.error_exit("phase_format must be a valid format string")
		self.formula_format = self.param_data.get('formula_format', self.formula_format)
		try:
			s = self.formula_format % u'str'
		except TypeError:
			self.error_exit("formula_format must be a valid format string")
		if self.param_data.get('current_phase') is not None:
			self.error_exit("current_phase mustn't be set, it is reserved for misp")


	######################################################################################################
	def generate_misptemp(self, current_phase):
		try:
			os.remove(self.misptemp_name)
		except OSError:
			pass
		self.param_data['current_phase'] = current_phase
		# TODO: this might be windows specific
		action_param_str = '"' + json.dumps(self.param_data).replace('"','\\"') + '"'
		retcode = self.execute_action_list([['', self.mispgen_action+' '+action_param_str]], self.display_mispgen)
		if retcode != 0:
			self.error_exit("misp generate action has failed")
		fileinfo = self.get_fileinfo(self.misptemp_name, self.DONT_USE_CACHE)
		if not fileinfo.exists:
			self.error_exit("misptemp '%s' is missing after generate action" % self.misptemp_name)
		first_pass = (len(current_phase) == 0)
		self.read_misptemp(first_pass)
		try:
			os.remove(self.misptemp_name)
		except OSError:
			pass
		self.merge_mispfile()


	######################################################################################################
	def write_misptemp(self, filename):
		misp_data = {}
		misp_data['options'] = self.options
		misp_data['configs'] = self.configs
		misp_data['phases'] = phases_data = []
		for phase in self.phases:
			phases_data.append(phase.to_json_data())
		ostr = json.dumps(misp_data, indent=self.JSON_INDENT, sort_keys=True)
		ofile = open(filename, 'wt')
		ofile.write(ostr)
		ofile.close()


	######################################################################################################
	def read_misptemp(self, first_pass):
		# decode json data
		try:
			ifile = open(self.misptemp_name)
			misp_data = json.load(ifile)
			ifile.close()
		except Exception as exc:
			self.error_exit('misptemp JSON load error: '+str(exc))
		# get options
		if first_pass:
			self.options = misp_data.get('options', self.options)
			self.multi_pass_gen = self.options.get('multi_pass_gen', self.multi_pass_gen)
		# get configs
		if first_pass:
			self.configs = misp_data.get('configs')
			if self.configs is None:
				self.error_exit('misptemp configs are missing')
			self.exec_phase_ids = self.configs.get(self.exec_config)
			if self.exec_phase_ids is None:
				self.error_exit("invalid exec_config '%s'"%self.exec_config)
			self.exec_phase_ids = set(self.exec_phase_ids)
		# get phases
		phases_data = misp_data.get('phases')
		if phases_data is None:
			self.error_exit('misptemp phases are missing')
		for phase_data in phases_data:
			id = phase_data.get('id')
			if id is None:
				self.error_exit('misptemp phase id is missing')
			if first_pass:
				self.add_phase(id, [], raise_exception=False)
			else:
				self.current_phase = None
				for phase in self.phases:
					if id == phase.id:
						self.current_phase = phase
						break
				if self.current_phase == None:
					self.error_exit("misptemp phase id '%s' was not generated in the first pass"%id)
			if not self.multi_pass_gen or not first_pass:
				formulas_data = phase_data.get('formulas')
				if formulas_data is None:
					self.error_exit('misptemp phase formulas are missing')
				if len(formulas_data) > 0:
					self.current_phase.newly_generated = True
				for id, formula_data in formulas_data.items():
					formula = self.add_formula(id, raise_exception=False)
					formula.from_json_data(formula_data)
					for action in formula.actions:
						if type(action)!=list or len(action)!=2 or type(action[0])!=type(u'') or type(action[1])!=type(u''):
							self.error_exit("misptemp formula '%s' has invalid action(s)"%id)


	######################################################################################################
	def write_mispfile(self):
		# if mispfile is not there then this is a 'clean' action
		fileinfo = self.get_fileinfo(self.mispfile_name, self.DONT_USE_CACHE)
		if not fileinfo.exists:
			return
		if len(self.options) == 0 and len(self.phases) == 0:
			return
		# convert data
		if self.multi_pass_gen:
			for prev_phase in self.prev_phases:
				phase = None
				for phase in self.phases:
					if phase.id == prev_phase.id:
						break
				if phase is None or phase.newly_generated:
					continue
				for id, prev_formula in prev_phase.formulas.items():
					prev_formula.newly_generated = False
					phase.formulas[id] = prev_formula
		misp_data = {}
		misp_data['options'] = self.options
		misp_data['configs'] = self.configs
		misp_data['phases'] = phases_data = []
		for phase in self.phases:
			phases_data.append(phase.to_json_data())
		# write file
		try:
			ofile = open(self.mispfile_name, 'wt')
		except IOError:
			self.error_exit("can't write mispfile '%s'" % (self.mispfile_name))
		json.dump(misp_data, ofile, indent=self.JSON_INDENT, sort_keys=True)
		ofile.close()


	######################################################################################################
	def read_mispfile(self):
		try:
			ifile = open(self.mispfile_name)
			istr = ifile.read()
			ifile.close()
		except IOError:
			ifile = open(self.mispfile_name, 'wt')
			ifile.close()
			return
		if len(istr) == 0:
			return
		try:
			misp_data = json.loads(istr)
		except Exception as exc:
			self.error_exit('mispfile JSON load error: '+str(exc))
		for phase_data in misp_data['phases']:
			id = phase_data.get('id')
			phase = misp_phase_t(id)
			self.prev_phases.append(phase)
			formulas_data = phase_data.get('formulas')
			for id, formula_data in formulas_data.items():
				formula = misp_formula_t(id)
				phase.formulas[id] = formula
				formula.from_json_data(formula_data)


	######################################################################################################
	def merge_mispfile(self):
		for prev_phase in self.prev_phases:
			keep_formulas = {}
			for id, prev_formula in prev_phase.formulas.items():
				formula = self.formulas.get(id)
				if formula == None:
					keep_formulas[id] = prev_formula
					continue
				formula.copy_state(prev_formula)
			prev_phase.formulas = keep_formulas


	######################################################################################################
	def execute_phase_list(self):
		for phase in self.phases:
			if phase.id not in self.exec_phase_ids:
				continue
			if self.display_phase == 'always':
				self.display_phase_id(phase)
			if self.multi_pass_gen:
				self.generate_misptemp(phase.id)
			self.fileinfo_map = {}
			# if exec_formula is specificed and it is in this phase then we execute only exec_formula and return immediately
			formula = self.formulas.get(self.exec_formula)
			if formula != None and formula.phase == phase:
				self.execute_formula(formula)
				break
			# otherwise execute all formulas in the phase
			else:
				for formula in phase.formulas.values():
					self.execute_formula(formula)


	######################################################################################################
	def display_phase_id(self, phase):
		if not phase.id_printed:
			sys.stdout.write(self.phase_format % phase.id + '\n')
			phase.id_printed = True


	######################################################################################################
	def execute_formula(self, formula):
		if self.display_formula == 'always':
			sys.stdout.write(self.formula_format % formula.id + '\n')
		# by default we don't need to do anything
		need_actions = False
		# check if formula is in state file
		if not formula.state_present:
			need_actions = True
		# check for empty target (clean actions)
		if not need_actions:
			if len(formula.targets)==0 and len(formula.extra_targets_name)==0:
				need_actions = True
		# check sources
		cur_sources_fileinfo = self.get_formula_fileinfo_list(formula, formula.sources, self.MUST_EXIST, self.USE_CACHE)
		if not need_actions:
			need_actions = self.check_fileinfo_list_changed(cur_sources_fileinfo, formula.sources_fileinfo, self.CHECK_EQUAL)
		formula.sources_fileinfo = cur_sources_fileinfo
		# check extra sources
		if not need_actions:
			if len(formula.extra_sources) > 0:
				cur_extra_sources_fileinfo = self.get_formula_fileinfo_list(formula, formula.extra_sources, self.DONT_HAVE_TO_EXIST, self.USE_CACHE)
				need_actions = self.check_fileinfo_list_changed(cur_extra_sources_fileinfo, formula.extra_sources_fileinfo, self.CHECK_EQUAL)
		# check targets
		if not need_actions:
			cur_targets_fileinfo = self.get_formula_fileinfo_list(formula, formula.targets, self.DONT_HAVE_TO_EXIST, self.DONT_USE_CACHE)
			need_actions = self.check_fileinfo_list_changed(cur_targets_fileinfo, formula.targets_fileinfo, self.CHECK_OLDER)
		# check extra targets
		if not need_actions:
			if len(formula.extra_targets) > 0:
				cur_extra_targets_fileinfo = self.get_formula_fileinfo_list(formula, formula.extra_targets, self.DONT_HAVE_TO_EXIST, self.USE_CACHE)
				need_actions = self.check_fileinfo_list_changed(cur_extra_targets_fileinfo, formula.extra_targets_fileinfo, self.CHECK_OLDER)
		# check actions agains previous actions
		if not need_actions:
			need_actions = self.check_action_list_changed(formula.prev_actions, formula.actions)
		# actions
		if need_actions:
			self.execute_formula_actions(formula)
			formula.targets_fileinfo = self.get_formula_fileinfo_list(formula, formula.targets, self.MUST_EXIST, self.DONT_USE_CACHE)
			if len(formula.extra_sources_name) > 0:
				formula.extra_sources = self.read_formula_extra_files(formula, formula.extra_sources_name)
				formula.extra_sources_fileinfo = self.get_formula_fileinfo_list(formula, formula.extra_sources, self.MUST_EXIST, self.USE_CACHE)
			if len(formula.extra_targets_name) > 0:
				formula.extra_targets = self.read_formula_extra_files(formula, formula.extra_targets_name)
				formula.extra_targets_fileinfo = self.get_formula_fileinfo_list(formula, formula.extra_targets, self.MUST_EXIST, self.DONT_USE_CACHE)


	######################################################################################################
	def get_formula_fileinfo_list(self, formula, name_list, must_exist, use_cache):
		place_str = "formula '%s'" % formula.id
		file_list = []
		for name in name_list:
			fileinfo = self.get_fileinfo(name, use_cache)
			if must_exist and not fileinfo.exists:
				self.error_exit("%s: file '%s' doesn't exist" % (place_str, name))
			file_list.append(fileinfo)
		return file_list


	######################################################################################################
	def check_fileinfo_list_changed(self, fileinfo_list_0, fileinfo_list_1, check_equal):
		if len(fileinfo_list_0) != len(fileinfo_list_1):
			return True
		for fileinfo_0 in fileinfo_list_0:
			changed = True
			for fileinfo_1 in fileinfo_list_1:
				if fileinfo_0.name == fileinfo_1.name:
					if check_equal:
						if fileinfo_0.mtime == fileinfo_1.mtime:
							changed = False
					else:
						if fileinfo_0.mtime >= fileinfo_1.mtime:
							changed = False
					break
			if changed:
				return True
		return False


	######################################################################################################
	def check_action_list_changed(self, prev_actions, actions):
		if len(prev_actions) != len(actions):
			return True
		for idx in range(len(prev_actions)):
			prev_action = prev_actions[idx]
			action = actions[idx]
			if prev_action[0] != action[0] or prev_action[1] != action[1]:
				return True
		return False


	######################################################################################################
	def execute_formula_actions(self, formula):
		if len(formula.actions) == 0:
			self.error_exit("formula '%s' has no actions to execute" % (formula.id))
		if self.display_phase == 'action':
			self.display_phase_id(formula.phase)
		if self.display_formula == 'action':
			sys.stdout.write(self.formula_format % formula.id + '\n')
		retcode = self.execute_action_list(formula.actions, True)
		if retcode != 0:
			formula.state_valid = False
			self.error_exit("formula '%s' actions returned with error code %d" % (formula.id, retcode))


	######################################################################################################
	def read_formula_extra_files(self, formula, name):
		try:
			extra_files_str = open(name, 'rt').read()
		except:
			formula.state_valid = False
			self.error_exit("formula '%s' can't read extra file '%s'" % (formula.id, name))
		extra_files_str = extra_files_str.replace('\r', '\n')
		extra_files_str = extra_files_str.replace('\\\n', '\n')
		extra_files_str = extra_files_str.replace('\n', '\t')
		name_list = []
		str_list = extra_files_str.split()
		for s in str_list:
			s = os.path.normpath(s)
			name_list.append(s)
		extra_files = []
		for name in name_list:
			if name[-1] != ':':
				extra_files.append(name)
		return extra_files


	######################################################################################################
	def get_fileinfo(self, file_name, use_cache):
		if use_cache:
			fileinfo = self.fileinfo_map.get(file_name)
			if fileinfo != None:
				return fileinfo
		fileinfo = misp_fileinfo_t(file_name, 0.0)
		try:
			file_stat = os.stat(file_name)
			fileinfo.exists = True
			if not stat.S_ISDIR(file_stat[stat.ST_MODE]):
				fileinfo.mtime = float(file_stat.st_mtime)
			else:
				fileinfo.mtime = 1.0
				fileinfo.is_dir = True
		except OSError:
			pass
		if use_cache:
			self.fileinfo_map[file_name] = fileinfo
		return fileinfo


	######################################################################################################
	def execute_action_list(self, actions, display_actions):
		retcode = 0
		for flag_action in actions:
			flags = flag_action[0]
			action = flag_action[1]
			flags_str = '('+flags+') ' if len(flags) > 0 else ''
			if display_actions:
				sys.stdout.write(flags_str+action+'\n')
			need_retcode = True
			if '-' in flags:
				need_retcode = False
			action_shell = 's' in flags
			try:
				proc = subprocess.Popen(action, shell=action_shell, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True)
				while True:
					output = proc.stdout.read(1)
					sys.stdout.write(output)
					retcode = proc.poll()
					if retcode != None:
						output = proc.stdout.read()
						sys.stdout.write(output)
						break
			except Exception as exc:
				retcode = -999
				sys.stdout.write(str(exc)+'\n')
			if not need_retcode:
				retcode = 0
			if retcode != 0:
				break
		return retcode


	######################################################################################################
	def set_option(self, key, value):
		self.options[key] = value


	######################################################################################################
	def add_config(self, id):
		if self.configs.get(id) is not None:
			self.error_exit("config id must be unique, '%s' is not" % id)
		self.configs[id] = []


	######################################################################################################
	def add_phase(self, id, configs, raise_exception=True):
		if len(id) == 0:
			self.error_exit("phase id must be non-empty", raise_exception)
		for phase in self.phases:
			if id == phase.id:
				self.error_exit("phase id must be unique, '%s' is already in use" % id, True)
		for config_name in configs:
			config = self.configs.get(config_name)
			if config is None:
				self.error_exit("invalid phase '%s' config '%s'" % (id,config_name), True)
			config.append(id)
		self.current_phase = misp_phase_t(id)
		self.phases.append(self.current_phase)
		return self.current_phase


	######################################################################################################
	def set_current_phase(self, id):
		if len(id) == 0:
			return
		self.current_phase = None
		for phase in self.phases:
			if id == phase.id:
				self.current_phase = phase
				break
		if self.current_phase == None:
			self.error_exit("invalid phase id '%s'" % id)


	######################################################################################################
	def add_formula(self, id, sources=[], targets=[], actions=[], extra_sources_name='',
					extra_targets_name='', raise_exception=True):
		if len(id) == 0:
			self.error_exit("formula id must be non-empty", raise_exception)
		formula = misp_formula_t(id, sources, targets, actions, extra_sources_name, extra_targets_name)
		if self.formulas.get(id) is not None:
			self.error_exit("formula id must be unique, '%s' is not" % id, raise_exception)
		formula.phase = self.current_phase
		self.current_phase.formulas[id] = formula
		self.formulas[id] = formula
		return formula


	######################################################################################################
	def error_exit(self, error, raise_exception=False):
		sys.stderr.write('MISP_ERROR: ' + error + '\n')
		self.write_mispfile()
		if raise_exception:
			raise Exception
		else:
			sys.exit(-1)


######################################################################################################################
if __name__ == '__main__':
	misp = misp_t()
	misp.main()
